# coding=utf-8
from numpy import *
import itertools
import time


def loadDataSet():  # 为了说明apriori的生成频繁项集
    dataSet = [[1, 2, 5],  # 1
               [2, 4],
               [2, 3],
               [1, 2, 4],  # 4
               [1, 3],
               [2, 3],
               [1, 3],
               [1, 2, 3, 5],  # 8
               [1, 2, 3]]
    return dataSet


def loadDataSet1():  # 说明生成的关联规则
    dataSet = [[1, 3, 4],
               [2, 3, 5],
               [1, 2, 3, 5],
               [1, 2, 3, 4, 5],
               [2, 5]]
    return dataSet


def loadDataSet2():  # 《机器学习实战》中的数据
    dataSet = [[1, 3, 4],
               [2, 3, 5],
               [1, 2, 3, 5],
               [2, 5]]
    return dataSet


# 创建候选项集C1
def createC1(dataSet):
    C1 = []
    for transaction in dataSet:
        for item in transaction:
            if [item] not in C1:
                # 每个元素的项集都用list类型
                C1.append([item])
    C1.sort()
    # 集合frozenset,不被用户修改，可以作为字典键值
    return map(frozenset, C1)


# 由候选集Ck，生成频繁项集Lk
def CkToLk(D, Ck, minSupport):
    # 记录候选集Ck中项集的出现的次数
    Ckcount = {}
    for transaction in D:
        for Ckitem in Ck:  # Ckitem是set型
            if Ckitem.issubset(transaction):
                # if Ckitem not in Ckcount.keys(): #非常耗时
                if not Ckcount.has_key(Ckitem):  # 快速
                    Ckcount[Ckitem] = 1
                else:
                    Ckcount[Ckitem] += 1
    # 总交易数
    numTransaction = float(len(D))  # 求支持度防止取整
    retList = []  # 存频繁k项集
    supportData = {}  # 存频繁k项集及对应的支持度
    # 计算支持度
    for key in Ckcount:
        support = Ckcount[key] / numTransaction
        if support >= minSupport:  # 去支持度小的项集
            # retList.append(key) #从后面添加
            retList.insert(0, key)  # 从前面添加
        supportData[key] = support
    return retList, supportData


# 由频繁k项集，生成候选k+1项集
def candidateGen(Lk, k):
    candidateCPrune = []  # 存剪枝后的候选k+1项集
    lenLk = len(Lk)  # 频繁k项集集合中的项集个数
    # 频繁k项集中，任取两个项集，
    # 若每个频繁k项集的前k-1个元素一样，则合并为候选k+1项集
    for i in range(lenLk):
        for j in range(i + 1, lenLk):
            L1 = list(Lk[i])[:k - 1]
            L2 = list(Lk[j])[:k - 1]
            L1.sort();
            L2.sort()  # 排序
            if L1 == L2:
                candidateC = Lk[i] | Lk[j]  # 合并
                addFlag = True  # 是否添加合并的候选集
                # ***裁剪，候选项集candidateC所含的k项子集，应该是频繁的
                # ***若不是频繁的，那么不要该候选项集
                # eg:[1,3,4]的2项子集有[1,3],[1,4],[3,4]
                for iter in itertools.combinations(candidateC, k):
                    # 用迭代器生成的是tuple类型，转换为frozenset
                    subCandidate = frozenset(list(iter))
                    # 判断候选项集的子集是否是频繁的
                    if subCandidate not in Lk:
                        addFlag = False
                        break
                if addFlag == True:
                    candidateCPrune.append(candidateC)
    return candidateCPrune


# 由apriori产生所有的频繁项集
def apriori(dataSet, minSupport=0.5):
    # 生成候选1项集：C1
    C1 = createC1(dataSet)
    print'C1=', C1
    # 将每条交易转成set集合类型
    D = map(set, dataSet)
    # 由C1->L1
    L1, supportData = CkToLk(D, C1, minSupport)
    print'L1=', L1
    print 'supportData1=', supportData
    # 存放所有的频繁项集
    L = [L1]
    k = 2
    # 由L1->C2->L2->C3...,
    # 索引从0开始，L[k-2]即为频繁k-1项集
    while (len(L[k - 2]) > 0):
        Ck = candidateGen(L[k - 2], k - 1)  # 频繁k-1项集生成候选k项集
        print'C%d=' % k, Ck
        # 由候选Ck生成频繁Lk
        Lk, supK = CkToLk(D, Ck, minSupport)
        print'L%d=' % k, Lk
        # 所有的候选项集及支持度数：
        supportData.update(supK)
        # 候选集Ck中所有元素对应的支持度
        print 'supportData%d=' % k, supK
        # 累计已经计算的支持度
        L.append(Lk)
        k += 1
    return L, supportData


# 产生关联规则:由L2,L3,....产生规则
# 输入：频繁项集列表L,支持度列表,最小置信度
# 输出：满足最小置信度的规则列表
def generateRules(L, supportData, minConf):
    # 存强关联规则
    strongRuleList = []

    # 从L2开始遍历
    for i in range(1, len(L)):
        # 遍历当前频繁项集的每一个项集
        # 如频繁2项集为：[[1,2]，[3,4]]
        for freqSet in L[i]:
            # 如取项集为[1,2]，则H1为[[1],[2]],单个元素组成
            H1 = [frozenset([item]) for item in freqSet]
            if i > 1:
                # 由频繁k项集，生成关联规则，其中k>2
                rulesFromConseq(freqSet, H1, supportData, strongRuleList, minConf)
            else:
                # 由频繁2项集得出的关联规则(直接根据置信度进行评估)
                calcConf(freqSet, H1, supportData, strongRuleList, minConf)
    return strongRuleList


# ***筛选满足最小置信度的关联规则
# freqSet：由产生freqSet来产生关联规则
# 候选的规则后件的集合H
# RuleList存强关联规则，实质为generateRules()中的strongRuleList
def calcConf(freqSet, H, supportData, RuleList, minConf):
    prunedH = []  # 满足最小置信度的关联规则后件的集合
    for conseq in H:
        # 计算置信度：freqSet-conseq -> conseq
        conf = supportData[freqSet] / supportData[freqSet - conseq]
        if conf >= minConf:
            prunedH.append(conseq)  # 添加满足最小置信度的关联规则后件
            # 强关联规则用箭头指向表示
            RuleList.append((freqSet - conseq, conseq, conf))
            # 打印
            print freqSet - conseq, '--->', conseq, ' conf:', conf
    return prunedH


# ***递归生成关联规则，改变RuleList即改变了strongRuleList
def rulesFromConseq(freqSet, H, supportData, RuleList, minConf):
    m = len(H[0])  # 候选的关联规则后件中项的长度

    if m == 1:  # 后件只有一个项
        # 过滤只有一个项的后件（产生规则如1,2,3->4，并返回过滤后的后件）
        H = calcConf(freqSet, H, supportData, RuleList, minConf)
    # 规则后件H的项的个数 比频繁项集freqSet中的项的个数少1，
    # 超过该条件无法产生关联规则
    if (len(freqSet) > (m + 1)):
        # 由后件H（m项集），生成候选的m+1项集,作为新的后件
        Hmp1 = candidateGen(H, m)
        # 由最小置信度过滤满足条件的关联规则后件(m+1项集)
        Hmp1 = calcConf(freqSet, Hmp1, supportData, RuleList, minConf)
        # 后件集合中不止一个，在后件集合中生成新的后件，产生新的关联规则
        # 递归
        if (len(Hmp1) > 1):
            rulesFromConseq(freqSet, Hmp1, supportData, RuleList, minConf)


# 频繁项集Lk-1，生成候选项集
if __name__ == '__main__':
    dataSet = loadDataSet()
L, supportdata = apriori(dataSet, minSupport=0.222)
print'L=', L

# #生成频繁项集
dataSet = loadDataSet1()
L, supportData = apriori(dataSet, minSupport=0.4)
print'L=', L
print'supportData=', supportData

# 产生关联规则
rules = generateRules(L, supportData, minConf=0.6)  # minConf=0.5看一下结果
print'rules=', rules

# ***毒蘑菇数据测试
mushDatSet = [line.split() for line in open('mushroom.dat').readlines()]
L, supportdata = apriori(mushDatSet, minSupport=0.3)
print'L=', L
for item in L[1]:  # 与毒蘑菇出现频繁的1项：

    if item.intersection('2'):
        print 'L2项集：', item
#
# #与毒蘑菇出现频繁的3项：
# for item in L[3]:
#     if item.intersection('2'):
#         print 'L4项集：',item
#